##################################################################
#
# PhyML Performance Visualizer
#
# Victor Hanson-Smith
# victorhs@cs.uoregon.edu
# August 2009 
#
# DESCRIPTION:
# This script parses the output from PhyML-M3l, captures lines with performance data,
# and then plots the data as PDF files. If the PhyML execution was configured with
# low verbosity, then the output will not contain performance data.  In this case,
# this script will exit gracefully and not create PDF plots.
# 
# USAGE:
# python plot_performance.py <filepath> <filepath> <filepath> etc.
#
# ... where <filepath> is the path to a text file capturing STDOUT from PhyML-M3l.  
#
# OUTPUT:
# Two CRAN scripts (i.e. 'R' scripts) will be written to the current working directory.  
# When these scripts are invoked, they will generate PDF plots.
# (1) plot1.cran plots iterations vs. -lnL
# (2) plot2.cran plots seconds vs. -lnL
#
#####################################################################

# rainbow
#colors = ["red", "blue", "orange", "purple", "green", "brown", "yellow", "aqua", "gold"]

# Unimax versus several Multimax
colors = ["orangered", "royalblue", "purple", "deepskyblue", "springgreen2", "gold2"]

#colors = ["orangered"]

for i in range(0,5):
    for j in range(0,5):
        colors.append( (i*j/1).__str__() )


import os
import sys
import re

filepath_points1 = {}
filepath_points2 = {}

for i in range(1, sys.argv.__len__()):
    filepath = None
    try:
        filepath = sys.argv[i]
    except IndexError:
        print "Yikes! I think you forgot to specify the filepath containing PhyML output."
        print "\nUsage: python plot_performance.py <filepath>"
    
    try:
        fin = open(filepath, "r")
    except IOError:
        print "Yikes! The file named ", filepath, " does not exist."
        exit()

    print "Opened file", filepath

    plot1_xy = {} # where key = iterations as integers, value = -lnL of the current tree at that iteration
    plot2_xy = {} # where key = milliseconds as doubles, value = current iteration at that time
    plot3_xy = {} # key = tree, value = lnL
    
    #
    # Parse the PhyML output, grab relevant data.
    #
    count_iterations = 0
    count_trees = 0
    last_hc_seconds = 0 # if we use a hyrbid approach (HC+STA), the clock gets reset after HC, so we need to adjust how we interpret the samples
    for line in fin.readlines():        
        if line.__contains__("sec") and line.startswith(". ("): # HC data
            line = line.strip()
            count_iterations += 1
            tokens = line.split()
            seconds_i = 2
            likelihood_i = 5
            if tokens[2].__contains__("sec"):
                seconds_i = 1
                likelihood_i = 4
            if line.__contains__("Topology"):
                count_trees += 1
            tokens[seconds_i] = re.sub("\(", "", tokens[seconds_i])
            tokens[seconds_i] = re.sub("\)", "", tokens[seconds_i])
            seconds = float(tokens[seconds_i])
            last_hc_seconds = seconds
            l = re.sub("\]", "", tokens[likelihood_i])
            likelihood = float(l)
            plot1_xy[ count_iterations ] = likelihood
            plot2_xy[ seconds ] = likelihood 
            
        elif line.__contains__("sec") and line.startswith("["): # STA data
            line = line.strip()
            count_iterations += 1
            tokens = line.split()
            s = tokens[0]
            s = re.sub("\[", "", s)
            seconds = float(s) + last_hc_seconds
            l = tokens[8]
            l = re.sub("\]", "", l)
            likelihood = float(l)
            
            if count_iterations % 10 == 0: # we'll skip some generations for plot1, else we'll suffer data explosion!
                plot1_xy[ count_iterations ] = likelihood
            plot2_xy[ seconds ] = likelihood
    fin.close()
        
    #
    # Sanity check
    #
    if plot1_xy.__len__() < 1:
        print "WARNING: I did not find any data for plot1 (iterations vs. -lnL) in file", filepath
    if plot2_xy.__len__() < 1:
        print "WARNING: I did not find any data for plot2 (time vs. iterations) in file", filepath
        
    filepath_points1[filepath] = plot1_xy
    filepath_points2[filepath] = plot2_xy

if filepath_points1.keys().__len__() < 1 or filepath_points2.keys().__len__() < 1:
    print "I found no data in one or more of your STDOUT traces."
    exit(1)


#
# Write CRAN scripts and plot PDFs.
#
def plot_in_r(points, output_filename_seed, title, xlab, ylab):
    maxx = None
    maxy = None
    miny = None

    filepaths = points.keys()
    filepaths.sort()

    # the ingredients for our CRAN script:
    plotstring = ""
    pointsstring = ""
    legendstring = "legend(\"bottomright\", c("
    for f in filepaths:
        legendstring += "\"" + f + "\","
    legendstring = re.sub(",$", "", legendstring)
    legendstring += "), cex=0.8, col=c("
    for col in colors:
        legendstring += "\"" + col + "\","
    legendstring = re.sub(",$", "", legendstring)        
    legendstring += "), pch=15);\n"

    count = -1
    for f in filepaths:
        #print points[f]
        count += 1
        x_sorted = points[f].keys()
        x_sorted.sort()
        
        string = "x" + count.__str__() + "<-c("
        for x in x_sorted:
            if maxx == None:
                maxx = x
            elif x > maxx:
                maxx = x
            string += x.__str__() + ","   
        string = re.sub(",$", "", string)
        string += ");\n"
        
        string += "y" + count.__str__() + "<-c("
        for x in x_sorted:
            if miny == None:
                miny = points[f][x]
            elif points[f][x] < miny:
                miny = points[f][x]
            if maxy == None:
                maxy = points[f][x]
            elif points[f][x] > maxy:
                maxy = points[f][x]
            string += points[f][x].__str__() + ","
        string = re.sub(",$", "", string)
        string += ");\n"
        string += "points(x" + count.__str__() + ",y" + count.__str__() + ", col='" + colors[count] + "', type='l',lwd=3);\n"
        pointsstring += string
    
    plotstring = "x <-c(0.0," + maxx.__str__() + ");\n"
    plotstring += "y <-c(" + miny.__str__() + "," + maxy.__str__() + ");\n"
    plotstring += "plot(x, y, type='n', main='" + title + "', xlab='" + xlab + "', ylab='" + ylab + "');\n"
    
    fout_cran = open(output_filename_seed + ".cran", "w")
    fout_cran.write("pdf('" + output_filename_seed + ".pdf', width=7, height=4);\n")
    fout_cran.write(plotstring)
    fout_cran.write(pointsstring)
    fout_cran.write(legendstring)
    fout_cran.write("dev.off();\n")
    fout_cran.close()
    
    os.system("r --no-save < " + output_filename_seed + ".cran")

plot_in_r(filepath_points1, "plot1", "-lnL(iterations)", "iterations", "-lnL")
plot_in_r(filepath_points2, "plot2", "-lnL(time)", "seconds", "-lnL")